#include <iostream>
#include <stdlib.h>
#include <cstring>
#include "acl/acl.h"
#include "acl/ops/acl_dvpp.h"

#define INFO_LOG(fmt, args...) fprintf(stdout, "[INFO] " fmt "\n", ##args)
#define WARN_LOG(fmt, args...) fprintf(stdout, "[WARN] " fmt "\n", ##args)
#define ERROR_LOG(fmt, args...) fprintf(stdout, "[ERROR] " fmt "\n", ##args)
#define ALIGN_UP(num, align) (((num) + ((align) - 1)) & ~((align) - 1))

using namespace std;

uint32_t deviceId = 0;
aclrtContext context = nullptr;
aclrtStream stream = nullptr;
acldvppChannelDesc * dvppChannelDesc = nullptr;
aclrtRunMode runMode;

aclError aclResourcesInit();    // 初始化acl环境
void aclResourcesFree();        // 释放acl资源
int ReadPicFile(string picName, void * &inputHostBuffer, int &inputDeviceBufferSize);
void writeVoidPtrToFile(const char * filename, void* data, size_t size);
int getAlignNum();

int main()
{
    string picPos = "../data/big.jpg";
    int width = 655;
    int height = 438;
    aclResourcesInit();

    // 准备输入数据
    void * inputHostBuffer;
    int inputDeviceBufferSize;
    int ret = ReadPicFile(picPos, inputHostBuffer, inputDeviceBufferSize);
    void * inputDeviceBuffer;
    aclError aclRet = acldvppMalloc(&inputDeviceBuffer, inputDeviceBufferSize);
    aclrtMemcpyKind kind;
    if (runMode == ACL_DEVICE) {
        kind = ACL_MEMCPY_DEVICE_TO_HOST;
    }
    else {
        kind = ACL_MEMCPY_HOST_TO_DEVICE;
    }
    aclRet = aclrtMemcpy(inputDeviceBuffer, inputDeviceBufferSize, inputHostBuffer, inputDeviceBufferSize, kind);

    // 预估JPEGD处理结果所需的内存大小
    uint32_t decodeOutBufferSize = 0;
    ret = acldvppJpegPredictDecSize(inputHostBuffer, inputDeviceBufferSize, PIXEL_FORMAT_YVU_SEMIPLANAR_420, &decodeOutBufferSize);
    free(inputHostBuffer);

    // 准备输出数据空间
    void * decodeOutDevBuffer;
    aclRet = acldvppMalloc(&decodeOutDevBuffer, decodeOutBufferSize);
    
    // 准备输出描述符
    acldvppPicDesc * decodeOutputDesc = acldvppCreatePicDesc();
    acldvppSetPicDescData(decodeOutputDesc, decodeOutDevBuffer);
    acldvppSetPicDescFormat(decodeOutputDesc, PIXEL_FORMAT_YUV_SEMIPLANAR_420); 
    acldvppSetPicDescWidth(decodeOutputDesc, ALIGN_UP(width, 2));
    acldvppSetPicDescHeight(decodeOutputDesc, ALIGN_UP(height, 2));
    acldvppSetPicDescWidthStride(decodeOutputDesc, ALIGN_UP(width, getAlignNum()));
    acldvppSetPicDescHeightStride(decodeOutputDesc, ALIGN_UP(height, 16));
    acldvppSetPicDescSize(decodeOutputDesc, decodeOutBufferSize);

    // 解码后销毁输出描述符
    aclRet = acldvppJpegDecodeAsync(dvppChannelDesc, inputDeviceBuffer, inputDeviceBufferSize, decodeOutputDesc, stream);
    ret = aclrtSynchronizeStream(stream);
    acldvppDestroyPicDesc(decodeOutputDesc);

    void * vpcOutHostBuffer = malloc(decodeOutBufferSize);
    if (runMode == ACL_DEVICE) {
        kind = ACL_MEMCPY_DEVICE_TO_HOST;
    }
    else {
        kind = ACL_MEMCPY_HOST_TO_DEVICE;
    }
    ret = aclrtMemcpy(vpcOutHostBuffer, decodeOutBufferSize, decodeOutDevBuffer, decodeOutBufferSize, kind);
    
    size_t slashPos = picPos.find_last_of('/');     // 找到最后一个 '/' 的位置
    size_t dotPos = picPos.find('.', slashPos);     // 找到第一个 '.' 的位置，从最后一个 '/' 之后开始找
    string outName = "../out/";
    outName = outName + to_string(ALIGN_UP(width, getAlignNum()));
    outName = outName + "x";
    outName = outName + to_string(ALIGN_UP(height, 16));
    outName = outName + "_";
    outName = outName + picPos.substr(slashPos + 1, dotPos - slashPos - 1);
    outName = outName + ".yuv";
    INFO_LOG("Output pic name \"%s\".", outName.c_str());
    writeVoidPtrToFile(outName.c_str(), vpcOutHostBuffer, decodeOutBufferSize);

    aclResourcesFree();
    return 0;
}


aclError aclResourcesInit()
{
    aclError ret = aclInit("");
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("Acl Init Failed, errorCode is %d", ret);
        return ret;
    }
    INFO_LOG("Acl Init Success");

    ret = aclrtSetDevice(deviceId);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("Acl Set Device Failed, errorCode is %d", ret);
        return ret;
    }
    INFO_LOG( "Acl Set Device Success, Current DeviceID:%d", deviceId);

    ret = aclrtCreateContext(&context, deviceId);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("Acl Create Context Failed, errorCode is %d", ret);
        return ret;
    }
    INFO_LOG("Acl Create Context Success, errorCode is %d", ret);

    ret = aclrtCreateStream(&stream);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("Acl Create Stream Failed, errorCode is %d", ret);
        return ret;
    }
    INFO_LOG("Acl Create Stream Success");

    dvppChannelDesc = acldvppCreateChannelDesc();
    if (dvppChannelDesc == nullptr) {
        ERROR_LOG("Acl acldvppCreateChannelDesc failed, errorCode is %d", ret);
        return ret;
    }
    INFO_LOG("Acl acldvppCreateChannelDesc Success");

    ret = acldvppCreateChannel(dvppChannelDesc);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("Acl acldvppCreateChannel failed, errorCode is %d", ret);
        return ret;
    }
    INFO_LOG("Acl acldvppCreateChannel Success");


    ret = aclrtGetRunMode(&runMode);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("Acl aclrtGetRunMode failed, errorCode is %d", ret);
        return ret;
    }
    INFO_LOG("Acl aclrtGetRunMode Success");


    return ret;
}

void aclResourcesFree()
{
    aclError ret;
    if (dvppChannelDesc != nullptr) {
        ret = acldvppDestroyChannel(dvppChannelDesc);
        if (ret != ACL_SUCCESS) {
            ERROR_LOG("Acl acldvppDestroyChannel failed, errorCode = %d", ret);
        }
        INFO_LOG("Acl AcldvppDestroyChannel Success");

        ret = acldvppDestroyChannelDesc(dvppChannelDesc);
        if (ret != ACL_SUCCESS) {
            ERROR_LOG("Acl acldvppDestroyChannelDesc failed, errorCode = %d", ret);
        }
        INFO_LOG("Acl acldvppDestroyChannelDesc Success");
        dvppChannelDesc = nullptr;
    }

    ret = aclrtDestroyStream(stream);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("Acl Destroy Stream Failed");
    }
    INFO_LOG("Acl Destroy Stream Success");

    ret = aclrtDestroyContext(context);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("Acl Destroy Context Failed");
    }
    INFO_LOG("Acl Destroy Context success");

    ret = aclrtResetDevice(deviceId);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("Acl Reset Device Failed");
    }
    INFO_LOG("Acl Reset Device Success");

    ret = aclFinalize();
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("Acl Finalize Failed");
    }
    INFO_LOG("Acl Finalize Success");
}

int ReadPicFile(string picName, void * &inputHostBuffer, int &inputDeviceBufferSize)
{
    FILE * fp = fopen(picName.c_str(), "rb");
    if (fp == NULL) {
        perror("Error opening file");
        return -1;
    }

    // 获取文件大小
    fseek(fp, 0, SEEK_END);
    inputDeviceBufferSize = ftell(fp);
    fseek(fp, 0, SEEK_SET);

    // 分配主机缓冲区内存
    inputHostBuffer = malloc(inputDeviceBufferSize);
    if (inputHostBuffer == NULL) {
        fclose(fp);
        perror("Error allocating memory");
        return -1;
    }

    // 读取文件内容到主机缓冲区
    size_t bytesRead = fread(inputHostBuffer, 1, inputDeviceBufferSize, fp);
    if (bytesRead!= static_cast<size_t>(inputDeviceBufferSize)) {
        fclose(fp);
        free(inputHostBuffer);
        perror("Error reading file");
        return -1;
    }

    fclose(fp);

    return 0;
}

void writeVoidPtrToFile(const char* filename, void* data, size_t size) 
{
    FILE* file = fopen(filename, "wb");
    if (file) {
        fwrite(data, size, 1, file);
        fclose(file);
    } else {
        perror("Error opening file for writing");
        exit(EXIT_FAILURE);
    }
}

int getAlignNum()
{
    const char * socVersion = aclrtGetSocName();
    int retNum = 128;
    if (strcmp(socVersion, "Ascend310B4") == 0) {
        retNum = 64;
    }
    else if (strcmp(socVersion, "Ascend310P3") == 0) {
        retNum = 64;
    }
    return retNum;
}